import base64
import json
import pickle
import time
import config
from utils import b64decode, b64encode, generate_timestamp_tuple
import numpy as np
import requests
from YOLOv3 import YOLOv3
import paddle
from img_loader import single_image_data_loader
from nms import nms
from draw_rec import save_results

ANCHORS = [10, 13, 16, 30, 33, 23, 30, 61, 62,
           45, 59, 119, 116, 90, 156, 198, 373, 326]
ANCHOR_MASKS = [[6, 7, 8], [3, 4, 5], [0, 1, 2]]
VALID_THRESH = 0.01
NMS_TOPK = 400
NMS_POSK = 100
NMS_THRESH = 0.45
NUM_CLASSES = 7
MODEL_STATE_DICT = paddle.load(r'../params/yolo_epoch50.pdparams')  # 模型参数

# 模型加载
model = YOLOv3(num_classes=NUM_CLASSES)
model.load_dict(MODEL_STATE_DICT)
model.eval()


# 加载图片，这一步只能在客户端做
def load_img(img_path: str):
    test_loader = single_image_data_loader(img_path, mode='test')
    data_generator = test_loader()  # test_loader()返回一个生成器
    img_name, img_data, img_scale_data = next(
        data_generator)  # img_data就是一个多维ndarray
    return img_name, img_data, img_scale_data


# 特征提取
def feature_extraction_local(last_result):
    """
    {
        "result": {}
        "timestamps": [("xxx", "local", 12345)]
    }
    """
    last_result['timestamps'].append(
        generate_timestamp_tuple("feature_extraction", "begin"))

    img_data = b64decode(last_result['result']['img_data'])

    img = paddle.to_tensor(img_data)
    outputs = model.forward(img)  # outputs是一个ndarray的list

    img_scale = last_result['result']['img_scale']
    last_result['result'] = {
        "outputs": b64encode(outputs),
        "img_scale": img_scale,
    }

    last_result['timestamps'].append(
        generate_timestamp_tuple("feature_extraction", "end"))
    return last_result


# 获取预测框以及得分 参数outputs是一个ndarray的list. 函数返回的都是ndarray的list
def get_bboxes_and_scores_local(last_result):

    last_result['timestamps'].append(generate_timestamp_tuple(
        "get_bboxes_and_scores", "begin"))
    # print(last_result['result'].keys())
    outputs = b64decode(last_result['result']['outputs'])
    img_scale = b64decode(last_result['result']['img_scale'])
    total_boxes = []
    total_scores = []

    img_scale = paddle.to_tensor(img_scale)
    downsample = 32

    for i, out in enumerate(outputs):
        out = paddle.to_tensor(out)
        anchor_mask = ANCHOR_MASKS[i]
        anchors_this_level = []
        for m in anchor_mask:
            anchors_this_level.append(ANCHORS[2 * m])
            anchors_this_level.append(ANCHORS[2 * m + 1])

        boxes, scores = paddle.vision.ops.yolo_box(
            x=out,
            img_size=img_scale,
            anchors=anchors_this_level,
            class_num=NUM_CLASSES,
            conf_thresh=VALID_THRESH,
            downsample_ratio=downsample,
            name="yolo_box" + str(i))

        scores = paddle.transpose(scores, perm=[0, 2, 1])
        total_boxes.append(boxes.numpy())
        total_scores.append(scores.numpy())
        downsample = downsample // 2

    last_result['result'] = {
        "total_boxes": b64encode(total_boxes),
        "total_scores": b64encode(total_scores)
    }
    last_result['timestamps'].append(
        generate_timestamp_tuple("get_bboxes_and_scores", "end"))
    return last_result


# 拼接不同层级的结果 total_boxes和total_scores都是ndarray的list. 函数的返回内容也都是ndarray的list
def concat_result_local(last_result):
    # total_boxes: list, total_scores: list
    last_result['timestamps'].append(
        generate_timestamp_tuple("concat_result", "begin"))
    total_boxes = b64decode(last_result['result']['total_boxes'])
    total_scores = b64decode(last_result['result']['total_scores'])

    for i in range(len(total_boxes)):
        total_boxes[i] = paddle.to_tensor(total_boxes[i])
    for i in range(len(total_scores)):
        total_scores[i] = paddle.to_tensor(total_scores[i])
    yolo_boxes = paddle.concat(total_boxes, axis=1)
    yolo_scores = paddle.concat(total_scores, axis=2)
    yolo_boxes = yolo_boxes.numpy()
    yolo_scores = yolo_scores.numpy()

    last_result['result'] = {
        "yolo_boxes": b64encode(yolo_boxes),
        "yolo_scores": b64encode(yolo_scores)
    }
    last_result['timestamps'].append(
        generate_timestamp_tuple("concat_result", "end"))
    return last_result

def multiclass_nms_local(last_result):
    return multiclass_nms(last_result,
                                  score_thresh=VALID_THRESH,
                                  nms_thresh=NMS_THRESH,
                                  pre_nms_topk=NMS_TOPK,
                                  pos_nms_topk=NMS_POSK)

# 多分类非极大值抑制
def multiclass_nms(last_result, score_thresh=0.01, nms_thresh=0.45, pre_nms_topk=1000, pos_nms_topk=100):

    last_result['timestamps'].append(
        generate_timestamp_tuple("multiclass_nms", "begin"))

    bboxes = b64decode(last_result['result']['yolo_boxes'])
    scores = b64decode(last_result['result']['yolo_scores'])

    batch_size = bboxes.shape[0]
    class_num = scores.shape[1]
    rets = []
    for i in range(batch_size):
        bboxes_i = bboxes[i]
        scores_i = scores[i]
        ret = []
        for c in range(class_num):
            scores_i_c = scores_i[c]
            keep_inds = nms(bboxes_i, scores_i_c, score_thresh,
                            nms_thresh, pre_nms_topk, i=i, c=c)
            if len(keep_inds) < 1:
                continue
            keep_bboxes = bboxes_i[keep_inds]
            keep_scores = scores_i_c[keep_inds]
            keep_results = np.zeros([keep_scores.shape[0], 6])
            keep_results[:, 0] = c
            keep_results[:, 1] = keep_scores[:]
            keep_results[:, 2:6] = keep_bboxes[:, :]
            ret.append(keep_results)
        if len(ret) < 1:
            rets.append(ret)
            continue
        ret_i = np.concatenate(ret, axis=0)
        scores_i = ret_i[:, 1]
        if len(scores_i) > pos_nms_topk:
            inds = np.argsort(scores_i)[::-1]
            inds = inds[:pos_nms_topk]
            ret_i = ret_i[inds]

        rets.append(ret_i)

    last_result['result'] = {
        'ret': b64encode(rets[0])
    }
    last_result['timestamps'].append(
        generate_timestamp_tuple("multiclass_nms", "end"))
    return last_result

